import warnings
from sklearn.exceptions import UndefinedMetricWarning
from train import TrainerReturnedObject

warnings.filterwarnings(action='ignore', category=UndefinedMetricWarning)
from loguru import logger

from sklearn.metrics import (confusion_matrix, accuracy_score, f1_score,
                             recall_score, precision_score,
                             classification_report)


class ClassificationEvaluator:
    def __init__(self, classes: list, class_to_idx: dict) -> None:
        """
        - param classes: 标签索引列表
        - param classes_to_idx: {标签索引:标签名}
        """
        self.classes = classes
        self.class_to_idx = class_to_idx
        self.trainer_returned_obj = TrainerReturnedObject()

    def update(self, obj: TrainerReturnedObject):
        self.trainer_returned_obj.update(obj)

    def loss(self):
        return self.trainer_returned_obj.losses / self.trainer_returned_obj.num if self.trainer_returned_obj.num > 0 else -1

    def accuracy(self):
        return accuracy_score(self.trainer_returned_obj.targets,
                              self.trainer_returned_obj.preds_argmax)

    def precision(self):
        return precision_score(self.trainer_returned_obj.targets,
                               self.trainer_returned_obj.preds_argmax,
                               average='macro')

    def recall(self):
        return recall_score(self.trainer_returned_obj.targets,
                            self.trainer_returned_obj.preds_argmax,
                            average='macro')

    def f1(self):
        return f1_score(self.trainer_returned_obj.targets,
                        self.trainer_returned_obj.preds_argmax,
                        average='macro')

    def report(self):
        return classification_report(self.trainer_returned_obj.targets,
                                     self.trainer_returned_obj.preds_argmax)

    def confusion_matrix(self):
        # return confusion_matrix(self.trainer_returned_obj.targets,
        #                         self.trainer_returned_obj.preds_argmax,
        #                         labels=self.class_to_idx)
        return confusion_matrix(self.trainer_returned_obj.targets,
                                self.trainer_returned_obj.preds_argmax)


class ClassificationLogger:
    def __init__(self, log_period=10) -> None:
        self.log_period = log_period

    def training_log(self, evaluator: ClassificationEvaluator, lr, step_now,
                     steps, epoch_now, epochs):
        if step_now % self.log_period == 0:
            logger.info(f'Epoch[{epoch_now}/{epochs}] '
                        f'Step[{step_now}/{steps}] '
                        f'lr {lr:.6f} '
                        f'loss {evaluator.loss():.6f} '
                        f'acc {evaluator.accuracy():.6f} '
                        f'f1 {evaluator.f1():.6f}')

    def validation_log(self, evaluator: ClassificationEvaluator, epoch_now,
                       epochs):
        logger.info(
            f'Epoch[{epoch_now}/{epochs}] loss {evaluator.loss():.6f}\n'
            f'Classification: \n'
            f'{evaluator.report()}')
